package utils

import (
	"bufio"
	"context"
	"fmt"
	"github.com/gorilla/websocket"
	"io"
	"net/http"
	"os/exec"
	"sync"
	"time"
)


type WebSocket struct {
	socket *websocket.Conn
	command *exec.Cmd
	wg       sync.WaitGroup
	cmdStr string
}

func InitWsConn(w http.ResponseWriter, r *http.Request, cmdStr string) (*WebSocket, error) {
	upgrader := websocket.Upgrader{
		ReadBufferSize:   1024,
		WriteBufferSize:  1024,
		HandshakeTimeout: 5 * time.Second,
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		return nil, err
	}
	return &WebSocket{
		socket: conn,
		wg: sync.WaitGroup{},
		cmdStr: cmdStr,
	}, nil
}


func (w *WebSocket) Push() {
	ctx, cancel := context.WithCancel(context.Background())
	fmt.Println("开始输出日志")

	c := exec.CommandContext(ctx, "bash", "-c", w.cmdStr)
	stdout, err := c.StdoutPipe()
	if err != nil {
		return
	}
	var wg sync.WaitGroup
	wg.Add(2)
	go w.readMsg(&wg, cancel)
	go w.writeMsg(&wg, stdout, ctx)
	err = c.Start()
	wg.Wait()
}

func (w *WebSocket) readMsg(wg *sync.WaitGroup, cancelFunc context.CancelFunc) {
	var data []byte
	var err error

	defer func() {
		fmt.Println("readMsg：减去一个 wg")
		wg.Done()
	}()

	for {
		if _, data, err = w.socket.ReadMessage(); err != nil {
			return
		}
		fmt.Println("读取的消息：", string(data))
		if string(data) == "END" {
			cancelFunc()
			return
		}
	}
}

func (w *WebSocket) writeMsg(wg *sync.WaitGroup, stdout io.ReadCloser, ctx context.Context) {
	defer func() {
		fmt.Println("writeMsg：减去一个 wg")
		wg.Done()
	}()

	reader := bufio.NewReader(stdout)
	for {
		// 其实这段去掉程序也会正常运行，只是我们就不知道到底什么时候Command被停止了，
		// 而且如果我们需要实时给web端展示输出的话，这里可以作为依据 取消展示
		select {
		// 检测到ctx.Done()之后停止读取
		case <-ctx.Done():
			if ctx.Err() != nil {
				fmt.Printf("程序出现错误: %q", ctx.Err())
			} else {
				fmt.Println("程序被终止")
			}
			return
		default:
			readString, err := reader.ReadString('\n')
			if err != nil || err == io.EOF {
				return
			}
			// fmt.Println("返回数据：", readString)
			if err := w.socket.WriteMessage(websocket.TextMessage, []byte(readString)); err != nil {
				fmt.Println("发送给客户端发送错误", err.Error())
			}
		}
	}
}
